@rickardp cc

@matthewdouglas if we can get this reviewed and merged, i can continue adding Metal support.

Hi @bghira

For the time being, I'd like to avoid linking in libtorch. Doing so adds more complexity to our packaging process, and IMO it is not worth it for only this. We aim to support a reasonably broad range of PyTorch versions; with linking to libtorch we'd have to either build for each of those and somehow distribute them all, or pin the version of PyTorch.

Eventually we'll consider using the newer LibTorch Stable ABI, but that would be quite a while away, and we're only going to do it if there's very clear benefit to doing so.

Instead, I ask that native code is built independently of PyTorch, and ideally exposes a C API that can be used in the same way that the CPU, CUDA, ROCm, and XPU backends work. Ideally we would not build as a Python extension that links to cpython either, but if so, I would also ask that it uses one build for all Python 3.10+ versions, e.g. with the Stable ABI.

I'm not super familiar here, but it seems the main things we need out of Torch are:

• torch::mps::get_dispatch_queue()
• torch::mps::get_command_buffer()
• torch::mps::commit()

It would be great to find a way around that. As mentioned on Discord, torch.mps.compile_shader() may also be an option. Do note that you can still write the shader in a separate file and read it in.

I understand that this is a little different then most PyTorch extensions with native code, but it's intentional to help us with distribution and broad compatibility.

understood; Metal support is basically not possible to take any further in this library without substantial performance drawbacks.